#define _CRT_SECURE_NO_WARNINGS
#include <stdio.h>
#include <process.h>
#include <dlmalloc.h>
#include "ThreadPool.h"

//------------------------------------------------------------------------------
struct Task 
{
    LPTHREAD_START_ROUTINE proc;
    void *data;
    struct list_head entry;
};
struct Thread 
{
    uint32_t id;
    int	is_free;
    HANDLE fd_wait;
    HANDLE fd_thread;
    struct list_head entry;
};

//------------------------------------------------------------------------------
ThreadPool::ThreadPool(int pool_size, int wait_for_dead_ms)
    : wait_for_dead_ms(wait_for_dead_ms),
      pool_size(pool_size), state(Destroyed)
{
    InitializeCriticalSection(&cs_task_queue);
    InitializeCriticalSection(&cs_thread_queue);

    INIT_LIST_HEAD(&task_queue);
    INIT_LIST_HEAD(&thread_queue);

    int ret = create();
}
ThreadPool::~ThreadPool()
{
    destroy();
    DeleteCriticalSection(&cs_task_queue);
    DeleteCriticalSection(&cs_thread_queue);
}

//------------------------------------------------------------------------------
u32 __stdcall ThreadPool::thread_proc(void *data)
{
    ThreadPool *pool = (ThreadPool *)data;
    if (!pool) return -1;

    HANDLE waits[2];
    u32 id = GetCurrentThreadId();
    struct Thread *t = pool->get_thread_from_queue(id);
    if (!t) return -1;

    waits[0] = t->fd_wait;
    waits[1] = pool->fd_shutdown;

    do {
        u32 e = WaitForMultipleObjects(2, waits, FALSE, INFINITE);
        if (e - WAIT_OBJECT_0 != 0)
            continue;

        struct Task *task = pool->get_task_from_queue();
        if (task && task->proc) 
            task->proc(task->data);

        pool->finish_notify(id);

    } while (!pool->check_state_stop());
    return 0;
}

//------------------------------------------------------------------------------
int ThreadPool::create()
{
    if (state != Destroyed) return -1;

    fd_shutdown = CreateEvent(NULL, TRUE, FALSE, NULL);
    if (!fd_shutdown) return -1;

    for (int i = 0; i < pool_size; i++)	{

        u32 id;
        HANDLE fd_thread = (HANDLE)_beginthreadex(NULL, 0, ThreadPool::thread_proc, 
			this, CREATE_SUSPENDED, (UINT *)&id);
        if (INVALID_HANDLE_VALUE == fd_thread)
            goto fail;

        struct Thread *thread = (Thread *)dlmalloc(sizeof(Thread));
        if (!thread) return -1;
        thread->is_free = 1;
        thread->fd_wait = CreateEvent(NULL, TRUE, FALSE, NULL);
        thread->fd_thread = fd_thread;
        thread->id = id;
        list_add_tail(&thread->entry, &thread_queue);
        ResumeThread(fd_thread);
    }
    state = Ready;
    return 0;

fail:
    struct list_head *tmp, *safe;
    list_for_each_safe(tmp, safe, &thread_queue) {
        struct Thread *thread = list_entry(tmp, struct Thread, entry);
        dlfree(thread);
        list_del(tmp);
    }
    return -1;
}

//------------------------------------------------------------------------------
void ThreadPool::destroy()
{
    struct list_head *tmp, *safe;

    if (state == Destroying || state == Destroyed)
        return;

    state = Destroying;

    if (get_working_thread_count() > 0)
        Sleep(wait_for_dead_ms);

    list_for_each_safe(tmp, safe, &thread_queue) {
        struct Thread *thread = list_entry(tmp, Thread, entry);
		CloseHandle(thread->fd_wait);
		CloseHandle(thread->fd_thread);
        dlfree(thread);
        list_del(tmp);
    }

	list_for_each_safe(tmp, safe, &task_queue) {
		struct Task *task = list_entry(tmp, Task, entry);
		dlfree(task);
		list_del(tmp);
	}

	CloseHandle(fd_shutdown);
    state = Destroyed;
}
//------------------------------------------------------------------------------
struct Task *ThreadPool::get_task_from_queue(void)
{
    struct Task *task;

    EnterCriticalSection(&cs_task_queue);
    task = list_first_entry_or_null(&task_queue, struct Task, entry);
    if (task) list_del(&task->entry);
    LeaveCriticalSection(&cs_task_queue);

    return task;
}
struct Thread *ThreadPool::get_thread_from_queue(u32 id)
{
	struct list_head *tmp;
    struct Thread *t = NULL;
    
    EnterCriticalSection(&cs_thread_queue);
    list_for_each(tmp, &thread_queue) {
        t = list_entry(tmp, Thread, entry);
        if (t->id == id) break;
    }
    LeaveCriticalSection(&cs_thread_queue);
    return t;
}

void ThreadPool::finish_notify(u32 id)
{
    struct Thread *t;
    struct list_head *tmp;

    EnterCriticalSection(&cs_thread_queue);
    list_for_each(tmp, &thread_queue) {
        t = list_entry(tmp, Thread, entry);
        if (t->id == id) {
            if (!list_empty(&task_queue)) goto out;
            t->is_free = 1;
            ResetEvent(t->fd_wait);
            break;
        };
    }
out:
    LeaveCriticalSection(&cs_thread_queue);
}

int ThreadPool::check_state_stop(void)
{
    int ret;

    EnterCriticalSection(&cs_task_queue);
    ret = (state == Destroying || state == Destroyed);
    LeaveCriticalSection(&cs_task_queue);

    return ret;
}
//------------------------------------------------------------------------------
int ThreadPool::run(LPTHREAD_START_ROUTINE proc, void *data,
    ThreadPriority priority)
{
    struct list_head *tmp;
    if (state == Destroying || state == Destroyed)
        return -1;

    struct Task *task = (struct Task *)dlmalloc(sizeof(struct Task));
    task->proc = proc;
    task->data = data;

    /* Put task into the task queue */
    EnterCriticalSection(&cs_task_queue);
    if (priority == Low)
        list_add_tail(&task->entry, &task_queue);
    else
        list_add(&task->entry, &task_queue);
    LeaveCriticalSection(&cs_task_queue);

    /* Find a free thread to do task */
    EnterCriticalSection(&cs_thread_queue);
    list_for_each(tmp, &thread_queue) {
        struct Thread *thread = list_entry(tmp, Thread, entry);

        if (thread->is_free) {
            thread->is_free = 0;
            SetEvent(thread->fd_wait);
            break;
        }
    }
    LeaveCriticalSection(&cs_thread_queue);

    return 0;
}

//------------------------------------------------------------------------------
int ThreadPool::get_working_thread_count()
{
    int cnt = 0;

    struct list_head *tmp;
    EnterCriticalSection(&cs_thread_queue);
    list_for_each(tmp, &thread_queue) {
        struct Thread *thread = list_entry(tmp, Thread, entry);

        if (!thread->is_free)
            cnt++;
    }
    LeaveCriticalSection(&cs_thread_queue);
    return cnt;
}
//------------------------------------------------------------------------------
PoolState ThreadPool::get_state()
{
    return state;
}
//------------------------------------------------------------------------------
